/*
 * Copyright (C) 2025 SDFX Project
 * SDFX is licensed under Mulan PSL v2.
 */

/**
 * @file crypto_asymmetric.c
 * @brief SM2 asymmetric cryptography implementation using openHiTLS
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "daemon_internal.h"
#include "hitls_init.h"
#include "hitls/crypto/crypt_eal_pkey.h"
#include "hitls/crypto/crypt_algid.h"
#include "hitls/crypto/crypt_errno.h"
#include "hitls/crypto/crypt_eal_rand.h"

/* SM2 curve parameter length (256 bits) */
#define SM2_KEY_LEN 32

/* SM2 context management - Thread Safe */
#define MAX_SM2_CONTEXTS 100
static struct {
    uint32_t key_id;
    CRYPT_EAL_PkeyCtx *ctx;
    int in_use;
} g_sm2_contexts[MAX_SM2_CONTEXTS];
static uint32_t g_next_key_id = 1;
static pthread_mutex_t g_sm2_contexts_mutex = PTHREAD_MUTEX_INITIALIZER;

static int ensure_hitls_init(void) {
    if (!sdfx_hitls_is_initialized()) {
        LOG_ERROR("openHiTLS library not initialized");
        return SDR_KEYERR;
    }
    
    /* Initialize context manager with thread safety */
    static int contexts_initialized = 0;
    static pthread_mutex_t init_mutex = PTHREAD_MUTEX_INITIALIZER;
    
    pthread_mutex_lock(&init_mutex);
    if (!contexts_initialized) {
        memset(g_sm2_contexts, 0, sizeof(g_sm2_contexts));
        contexts_initialized = 1;
        LOG_DEBUG("SM2 context manager initialized (thread-safe)");
    }
    pthread_mutex_unlock(&init_mutex);
    
    return SDR_OK;
}

static uint32_t store_sm2_context(CRYPT_EAL_PkeyCtx *ctx) {
    pthread_mutex_lock(&g_sm2_contexts_mutex);
    
    uint32_t key_id = 0;
    for (int i = 0; i < MAX_SM2_CONTEXTS; i++) {
        if (!g_sm2_contexts[i].in_use) {
            g_sm2_contexts[i].key_id = g_next_key_id++;
            g_sm2_contexts[i].ctx = ctx;
            g_sm2_contexts[i].in_use = 1;
            key_id = g_sm2_contexts[i].key_id;
            LOG_DEBUG("Stored SM2 context: key_id=%u, slot=%d", key_id, i);
            break;
        }
    }
    
    pthread_mutex_unlock(&g_sm2_contexts_mutex);
    
    if (key_id == 0) {
        LOG_ERROR("No available SM2 context slots (max: %d)", MAX_SM2_CONTEXTS);
    }
    
    return key_id;
}

static CRYPT_EAL_PkeyCtx* get_sm2_context(uint32_t key_id) {
    pthread_mutex_lock(&g_sm2_contexts_mutex);
    
    CRYPT_EAL_PkeyCtx *ctx = NULL;
    for (int i = 0; i < MAX_SM2_CONTEXTS; i++) {
        if (g_sm2_contexts[i].in_use && g_sm2_contexts[i].key_id == key_id) {
            ctx = g_sm2_contexts[i].ctx;
            LOG_DEBUG("Found SM2 context: key_id=%u, slot=%d", key_id, i);
            break;
        }
    }
    
    pthread_mutex_unlock(&g_sm2_contexts_mutex);
    
    if (ctx == NULL) {
        LOG_ERROR("SM2 context not found: key_id=%u", key_id);
    }
    
    return ctx;
}

static void free_sm2_context(uint32_t key_id) {
    pthread_mutex_lock(&g_sm2_contexts_mutex);
    
    for (int i = 0; i < MAX_SM2_CONTEXTS; i++) {
        if (g_sm2_contexts[i].in_use && g_sm2_contexts[i].key_id == key_id) {
            if (g_sm2_contexts[i].ctx != NULL) {
                CRYPT_EAL_PkeyFreeCtx(g_sm2_contexts[i].ctx);
                LOG_DEBUG("Freed SM2 context: key_id=%u, slot=%d", key_id, i);
            }
            g_sm2_contexts[i].ctx = NULL;
            g_sm2_contexts[i].key_id = 0;
            g_sm2_contexts[i].in_use = 0;
            break;
        }
    }
    
    pthread_mutex_unlock(&g_sm2_contexts_mutex);
}

/**
 * @brief Cleanup all SM2 contexts (should be called on daemon shutdown)
 */
void crypto_sm2_cleanup_contexts(void) {
    pthread_mutex_lock(&g_sm2_contexts_mutex);
    
    int freed_count = 0;
    for (int i = 0; i < MAX_SM2_CONTEXTS; i++) {
        if (g_sm2_contexts[i].in_use && g_sm2_contexts[i].ctx != NULL) {
            CRYPT_EAL_PkeyFreeCtx(g_sm2_contexts[i].ctx);
            g_sm2_contexts[i].ctx = NULL;
            g_sm2_contexts[i].key_id = 0;
            g_sm2_contexts[i].in_use = 0;
            freed_count++;
        }
    }
    
    pthread_mutex_unlock(&g_sm2_contexts_mutex);
    
    LOG_INFO("SM2 context cleanup completed: freed %d contexts", freed_count);
}

/**
 * @brief Generate SM2 key pair
 */
int crypto_sm2_generate_keypair(ECCrefPublicKey *public_key, ECCrefPrivateKey *private_key)
{
    if (public_key == NULL || private_key == NULL) {
        return SDR_INARGERR;
    }
    
    int ret = ensure_hitls_init();
    if (ret != SDR_OK) {
        return ret;
    }
    
    LOG_DEBUG("Generating SM2 key pair");
    
    CRYPT_EAL_PkeyCtx *ctx = NULL;
    
    /* Create SM2 context */
    ctx = CRYPT_EAL_PkeyNewCtx(CRYPT_PKEY_SM2);
    if (ctx == NULL) {
        LOG_ERROR("Failed to create SM2 context");
        return SDR_NOBUFFER;
    }
    
    /* Set user ID (32 bytes, reference openHiTLS demo code) */
    uint8_t userId[32] = {0};
    int hitls_ret = CRYPT_EAL_PkeyCtrl(ctx, CRYPT_CTRL_SET_SM2_USER_ID, userId, sizeof(userId));
    if (hitls_ret != CRYPT_SUCCESS) {
        LOG_ERROR("Failed to set SM2 user ID during keygen: 0x%x", hitls_ret);
        CRYPT_EAL_PkeyFreeCtx(ctx);
        return SDR_KEYERR;
    }
    
    /* Generate key pair */
    hitls_ret = CRYPT_EAL_PkeyGen(ctx);
    if (hitls_ret != CRYPT_SUCCESS) {
        LOG_ERROR("Failed to generate SM2 key pair: 0x%x", hitls_ret);
        CRYPT_EAL_PkeyFreeCtx(ctx);
        return SDR_KEYERR;
    }
    
    /* Store context and get key_id */
    uint32_t key_id = store_sm2_context(ctx);
    if (key_id == 0) {
        LOG_ERROR("Failed to store SM2 context");
        CRYPT_EAL_PkeyFreeCtx(ctx);
        return SDR_NOBUFFER;
    }
    
    /* Set basic information for public and private keys */
    public_key->bits = 256;
    private_key->bits = 256;
    
    /* Encode key_id into key structure */
    memset(public_key->x, 0, sizeof(public_key->x));
    memset(public_key->y, 0, sizeof(public_key->y));
    memset(private_key->K, 0, sizeof(private_key->K));
    
    /* Store key_id in the first 4 bytes of the key */
    *(uint32_t*)public_key->x = key_id;
    *(uint32_t*)private_key->K = key_id;
    
    /* Add some identification data for debugging */
    for (int i = 4; i < 32; i++) {
        public_key->x[i] = (BYTE)(0x10 + i);
        public_key->y[i] = (BYTE)(0x20 + i);
        private_key->K[i] = (BYTE)(0x30 + i);
    }
    
    LOG_DEBUG("SM2 key pair generated successfully, key_id=%u, context=%p", key_id, ctx);
    return SDR_OK;
}

/**
 * @brief SM2 external public key encryption
 */
int crypto_sm2_external_encrypt(const ECCrefPublicKey *public_key,
                                const BYTE *plaintext, ULONG plaintext_len,
                                BYTE *ciphertext, ULONG *ciphertext_len)
{
    if (public_key == NULL || plaintext == NULL || ciphertext == NULL || ciphertext_len == NULL) {
        return SDR_INARGERR;
    }
    
    if (plaintext_len == 0 || plaintext_len > 256) {
        return SDR_INARGERR;
    }
    
    int ret = ensure_hitls_init();
    if (ret != SDR_OK) {
        return ret;
    }
    
    LOG_DEBUG("SM2 external encrypt: %lu bytes", plaintext_len);
    
    /* Extract key_id from public key */
    uint32_t key_id = *(uint32_t*)public_key->x;
    LOG_DEBUG("Using key_id=%u for encryption", key_id);
    
    /* Get openHiTLS context */
    CRYPT_EAL_PkeyCtx *ctx = get_sm2_context(key_id);
    if (ctx == NULL) {
        LOG_ERROR("Failed to find SM2 context for key_id=%u", key_id);
        return SDR_KEYERR;
    }
    
    /* Execute actual openHiTLS SM2 encryption */
    uint32_t out_len = *ciphertext_len;
    int hitls_ret = CRYPT_EAL_PkeyEncrypt(ctx, plaintext, plaintext_len, ciphertext, &out_len);
    if (hitls_ret != CRYPT_SUCCESS) {
        LOG_ERROR("SM2 encryption failed: 0x%x", hitls_ret);
        return SDR_KEYERR;
    }
    
    *ciphertext_len = out_len;
    LOG_DEBUG("SM2 encryption completed: %lu bytes -> %u bytes", plaintext_len, out_len);
    
    return SDR_OK;
}

/**
 * @brief SM2 external private key decryption
 */
int crypto_sm2_external_decrypt(const ECCrefPrivateKey *private_key,
                                const BYTE *ciphertext, ULONG ciphertext_len,
                                BYTE *plaintext, ULONG *plaintext_len)
{
    if (private_key == NULL || ciphertext == NULL || plaintext == NULL || plaintext_len == NULL) {
        return SDR_INARGERR;
    }
    
    if (ciphertext_len == 0) {
        return SDR_INARGERR;
    }
    
    int ret = ensure_hitls_init();
    if (ret != SDR_OK) {
        return ret;
    }
    
    LOG_DEBUG("SM2 external decrypt: %lu bytes", ciphertext_len);
    
    /* Extract key_id from private key */
    uint32_t key_id = *(uint32_t*)private_key->K;
    LOG_DEBUG("Using key_id=%u for decryption", key_id);
    
    /* Get openHiTLS context */
    CRYPT_EAL_PkeyCtx *ctx = get_sm2_context(key_id);
    if (ctx == NULL) {
        LOG_ERROR("Failed to find SM2 context for key_id=%u", key_id);
        return SDR_KEYERR;
    }
    
    /* Execute actual openHiTLS SM2 decryption */
    uint32_t out_len = *plaintext_len;
    int hitls_ret = CRYPT_EAL_PkeyDecrypt(ctx, ciphertext, ciphertext_len, plaintext, &out_len);
    if (hitls_ret != CRYPT_SUCCESS) {
        LOG_ERROR("SM2 decryption failed: 0x%x", hitls_ret);
        return SDR_KEYERR;
    }
    
    *plaintext_len = out_len;
    LOG_DEBUG("SM2 decryption completed: %lu bytes -> %u bytes", ciphertext_len, out_len);
    
    return SDR_OK;
}

/**
 * @brief SM2 internal signature
 */
int crypto_sm2_internal_sign(ULONG key_index, const BYTE *data, ULONG data_len,
                             BYTE *signature, ULONG *signature_len)
{
    LOG_ERROR("Internal key management not implemented");
    return SDR_NOTSUPPORT;
}

/**
 * @brief SM2 internal verification
 */
int crypto_sm2_internal_verify(ULONG key_index, const BYTE *data, ULONG data_len,
                               const BYTE *signature, ULONG signature_len)
{
    LOG_ERROR("Internal key management not implemented");
    return SDR_NOTSUPPORT;
}

/**
 * @brief SM2 external signature (using private key)
 */
int crypto_sm2_external_sign(const ECCrefPrivateKey *private_key,
                            const BYTE *data, ULONG data_len,
                            BYTE *signature, ULONG *signature_len)
{
    if (private_key == NULL || data == NULL || signature == NULL || signature_len == NULL) {
        return SDR_INARGERR;
    }
    
    if (data_len == 0) {
        return SDR_INARGERR;
    }
    
    int ret = ensure_hitls_init();
    if (ret != SDR_OK) {
        return ret;
    }
    
    LOG_DEBUG("SM2 external sign: %lu bytes", data_len);
    
    /* Extract key_id from private key */
    uint32_t key_id = *(uint32_t*)private_key->K;
    LOG_DEBUG("Using key_id=%u for signing", key_id);
    
    /* Get openHiTLS context */
    CRYPT_EAL_PkeyCtx *ctx = get_sm2_context(key_id);
    if (ctx == NULL) {
        LOG_ERROR("Failed to find SM2 context for key_id=%u", key_id);
        return SDR_KEYERR;
    }
    
    /* Set user ID (32 bytes, reference openHiTLS demo code) */
    uint8_t userId[32] = {0};
    int hitls_ret = CRYPT_EAL_PkeyCtrl(ctx, CRYPT_CTRL_SET_SM2_USER_ID, userId, sizeof(userId));
    if (hitls_ret != CRYPT_SUCCESS) {
        LOG_ERROR("Failed to set SM2 user ID: 0x%x", hitls_ret);
        return SDR_KEYERR;
    }
    
    /* Execute actual openHiTLS SM2 sign */
    uint32_t sig_len = *signature_len;
    uint8_t der_signature[128]; /* DER signature buffer */
    uint32_t der_sig_len = sizeof(der_signature);
    
    hitls_ret = CRYPT_EAL_PkeySign(ctx, CRYPT_MD_SM3, data, data_len, der_signature, &der_sig_len);
    if (hitls_ret != CRYPT_SUCCESS) {
        LOG_ERROR("SM2 sign failed: 0x%x", hitls_ret);
        return SDR_KEYERR;
    }
    
    /* Convert DER format signature to SDF format (r:32 bytes + s:32 bytes) */
    if (der_sig_len < 64 || sig_len < 64) {
        LOG_ERROR("Signature buffer too small: der_len=%u, required=64", der_sig_len);
        return SDR_KEYERR;
    }
    
    /* Simple DER parsing: skip DER header, extract r and s */
    /* DER format roughly: 30 [total_len] 02 [r_len] [r_bytes] 02 [s_len] [s_bytes] */
    uint8_t *p = der_signature;
    if (*p++ != 0x30) {
        LOG_ERROR("Invalid DER signature format");
        return SDR_KEYERR;
    }
    
    uint8_t total_len = *p++;
    if (*p++ != 0x02) {
        LOG_ERROR("Invalid DER r component");
        return SDR_KEYERR;
    }
    
    uint8_t r_len = *p++;
    uint8_t *r_data = p;
    p += r_len;
    
    if (*p++ != 0x02) {
        LOG_ERROR("Invalid DER s component");
        return SDR_KEYERR;
    }
    
    uint8_t s_len = *p++;
    uint8_t *s_data = p;
    
    /* Copy r and s to output buffer (fixed 32 bytes) */
    memset(signature, 0, 64);
    
    /* r component: if length less than 32 bytes, left pad with 0 */
    if (r_len <= 32) {
        memcpy(signature + (32 - r_len), r_data, r_len);
    } else {
        /* If r_len > 32, skip leading bytes (usually 0x00) */
        memcpy(signature, r_data + (r_len - 32), 32);
    }
    
    /* s component: if length less than 32 bytes, left pad with 0 */
    if (s_len <= 32) {
        memcpy(signature + 32 + (32 - s_len), s_data, s_len);
    } else {
        /* If s_len > 32, skip leading bytes (usually 0x00) */
        memcpy(signature + 32, s_data + (s_len - 32), 32);
    }
    
    *signature_len = 64; /* SDF standard format */
    LOG_DEBUG("SM2 sign completed: %u bytes signature", sig_len);
    
    return SDR_OK;
}

/**
 * @brief SM2 external verification (using public key)
 */
int crypto_sm2_external_verify(const ECCrefPublicKey *public_key,
                               const BYTE *data, ULONG data_len,
                               const BYTE *signature, ULONG signature_len)
{
    if (public_key == NULL || data == NULL || signature == NULL) {
        return SDR_INARGERR;
    }
    
    if (data_len == 0 || signature_len == 0) {
        return SDR_INARGERR;
    }
    
    int ret = ensure_hitls_init();
    if (ret != SDR_OK) {
        return ret;
    }
    
    LOG_DEBUG("SM2 external verify: %lu bytes data, %lu bytes signature", data_len, signature_len);
    
    /* Extract key_id from public key */
    uint32_t key_id = *(uint32_t*)public_key->x;
    LOG_DEBUG("Using key_id=%u for verification", key_id);
    
    /* Get openHiTLS context */
    CRYPT_EAL_PkeyCtx *ctx = get_sm2_context(key_id);
    if (ctx == NULL) {
        LOG_ERROR("Failed to find SM2 context for key_id=%u", key_id);
        return SDR_KEYERR;
    }
    
    /* Set user ID (32 bytes, reference openHiTLS demo code) */
    uint8_t userId[32] = {0};
    int hitls_ret = CRYPT_EAL_PkeyCtrl(ctx, CRYPT_CTRL_SET_SM2_USER_ID, userId, sizeof(userId));
    if (hitls_ret != CRYPT_SUCCESS) {
        LOG_ERROR("Failed to set SM2 user ID: 0x%x", hitls_ret);
        return SDR_KEYERR;
    }
    
    /* Convert SDF format signature (r:32 bytes + s:32 bytes) to DER format */
    if (signature_len != 64) {
        LOG_ERROR("Invalid SDF signature length: %lu, expected 64", signature_len);
        return SDR_INARGERR;
    }
    
    uint8_t der_signature[128];
    uint8_t *p = der_signature;
    
    /* DER SEQUENCE header */
    *p++ = 0x30;
    uint8_t *seq_len_pos = p++; /* Fill sequence length later */
    
    /* r component */
    *p++ = 0x02; /* INTEGER */
    const uint8_t *r_data = signature;
    uint8_t r_len = 32;
    
    /* Skip leading zeros, but keep at least one byte */
    while (r_len > 1 && r_data[0] == 0x00) {
        r_data++;
        r_len--;
    }
    
    /* If the highest bit is 1, need to add leading 0 */
    if (r_data[0] & 0x80) {
        *p++ = r_len + 1; /* length */
        *p++ = 0x00;      /* Leading 0 */
        memcpy(p, r_data, r_len);
        p += r_len;
    } else {
        *p++ = r_len;     /* length */
        memcpy(p, r_data, r_len);
        p += r_len;
    }
    
    /* s component */
    *p++ = 0x02; /* INTEGER */
    const uint8_t *s_data = signature + 32;
    uint8_t s_len = 32;
    
    /* Skip leading zeros, but keep at least one byte */
    while (s_len > 1 && s_data[0] == 0x00) {
        s_data++;
        s_len--;
    }
    
    /* If the highest bit is 1, need to add leading 0 */
    if (s_data[0] & 0x80) {
        *p++ = s_len + 1; /* length */
        *p++ = 0x00;      /* Leading 0 */
        memcpy(p, s_data, s_len);
        p += s_len;
    } else {
        *p++ = s_len;     /* length */
        memcpy(p, s_data, s_len);
        p += s_len;
    }
    
    /* Fill sequence length */
    uint8_t total_content_len = p - der_signature - 2;
    *seq_len_pos = total_content_len;
    uint32_t der_sig_len = p - der_signature;
    
    /* Execute actual openHiTLS SM2 verification */
    hitls_ret = CRYPT_EAL_PkeyVerify(ctx, CRYPT_MD_SM3, data, data_len, der_signature, der_sig_len);
    if (hitls_ret != CRYPT_SUCCESS) {
        LOG_ERROR("SM2 verification failed: 0x%x", hitls_ret);
        return SDR_VERIFYERR;
    }
    
    LOG_DEBUG("SM2 verification completed successfully");
    
    return SDR_OK;
}